/* Copyright (C) 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012 Mario Orsi
   This file is part of Brahms.
   Brahms is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
   Brahms is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more details.
   You should have received a copy of the GNU General Public License along with Brahms.  If not, see <http://www.gnu.org/licenses/>. */

/****************************************************************************************************************
 * output.c -- functions for outputting (averages of) data of interest and coordinetates' snapshot/trajectories * 
 ****************************************************************************************************************/

#include "dataStructs.h"

extern const VecR region, *siteCoords, *siteOrientVec; 
extern const RMat *siteRotMat;
extern const OrdParSolute ordParSolute;
extern const Site *site;
extern const Atom *atom;
extern const LipidStruct dopcStruct;
extern const VecRTwo *xyLipidComTrue, xySoluteComTrue;
extern const real timeNow;
extern const Prop totEnergy, potEnergy, systemTemperature, lipTemperature, watTemperature, solTemperature, pressure, dens;
extern const Prop boxVolume, xyArea, lipVolume, lipArea, surfaceTension, headGroupDipole, headGroupTilt;
extern const Prop ordParTail, ordParTop, ordParMidUnsat, ordParMidSat, ordParEnd, zConstraintForce;
extern const int nSites, nLipids, nSolutes, stepCount;
extern FILE *fPtrZCF, *fPtrAV;

/* external function declarations: */
void RotMatToEuler( real*, real*, real*, RMat* ); 
void RotMatToQuat( Quat*, RMat );
void ErrExit( int );

void PrintBrahmsVersion()
{
  printf("\n**********************\n"); 
  printf(  "* BRAHMS - 09Feb2012 *\n"); 
  printf(  "**********************\n\n");
  fflush(stdout);
}

void PrintMatrix( RMat M ) 
{
  printf( "% 8.3f % 8.3f % 8.3f\n",   M.u[0], M.u[3], M.u[6] );
  printf( "% 8.3f % 8.3f % 8.3f\n",   M.u[1], M.u[4], M.u[7] );
  printf( "% 8.3f % 8.3f % 8.3f\n\n", M.u[2], M.u[5], M.u[8] );
}

void PrintHeader( FILE *fp ) 
{
  if ( nLipids ) {
    fprintf( fp, "\n" );
    fprintf( fp, "   step    time/ns  " ); 
    fprintf( fp, "E/(kJ/mol)  U/(kJ/mol)  P/atm  T_lip/C  T_H2O/C  " ); 
    if ( nSolutes ) fprintf( fp, "%10s", "T-sol/C" );
    fprintf( fp, "A_lip/(A^2)  V_lip/(nm^3)  ST/(dyn/cm)  dipHG/D  tiltHG/deg" );
    //    fprintf( fp, "%7s %7s %7s %8s %6s", "P2Top", "P2Mid", "P2Bot", "P2VBot", "Smol" );
    /*     fprintf( fp, " %4s", "Smol" ); */
    //fprintf( "%11s %10s ", "GlyTiltUp", "GlyTiltLo" );
    fprintf( fp, "\n" );
  } else { // if waters only
    fprintf( fp, "%8s %9s %18s %18s", "step", "time/ns", "Etot/(kJ/mol) sd", "Epot/(kJ/mol) sd" );
    fprintf( fp, "%14s %14s", "P/atm   sd", "T/C   sd" ); 
    //    fprintf( fp, "%7s %5s", "wT", "sd(wT)" ); 
    // fprintf( fp, "%16s %16s", "V [nm^3] sd(V)", "dens [g/cm^3] sd(dens)" ); 
    fprintf( fp, "%20s", "dens/(g/cm^3) sd" ); 
    //    fprintf( fp, "%20s", "Cv/(J/(mol*K))" ); 
    fprintf( fp, "\n");
  }
  fflush(fp);
}

void PrintZConstraintHeader() {
  FILE *fPtr;
  fPtr = fopen( "zConstraint.out", "a" );
  fprintf( fPtr, "Time/ns   fzc[kJ/(mol*nm)]  sd(fzc)" );
  fprintf( fPtr, "\n" );
  fclose( fPtr );
}

void PrintAvgZConstraintParameters() {
  FILE *fPtr;
  fPtr = fopen( "zConstraint.out", "a" );
  fprintf( fPtr, "%6.1f\t", timeNow * TIME_ns );
  fprintf( fPtr, "% 12.2f   %12.1f\t", zConstraintForce.sum * FORCE_kJ_mol_nm, zConstraintForce.sum2 * FORCE_kJ_mol_nm );
  fprintf( fPtr, "\n" );
  fclose( fPtr );
  if ( ( fPtr = fopen( "ordParSolute.dat", "a" ) ) == NULL ) printf( "File ordParSolute.dat could not be created\n" );
  else { 
    fprintf( fPtr, "% 5.4f % 5.4f\t\t", ordParSolute.p2x.sum, ordParSolute.p2x.sum2 );
    fprintf( fPtr, "% 5.4f % 5.4f\t\t", ordParSolute.p2y.sum, ordParSolute.p2y.sum2 );
    fprintf( fPtr, "% 5.4f % 5.4f",     ordParSolute.p2z.sum, ordParSolute.p2z.sum2 );
    fprintf( fPtr, "\n" );
    fclose( fPtr ); 
  }  
}

void PrintZConstraintForce() 
{
  fprintf( fPtrZCF, "% 9.3f\n", zConstraintForce.val * FORCE_kJ_mol_nm );
}

void PrintWatLatMotion() // used when z-constraining 
{
  FILE *fPtr;
  fPtr = fopen( "xyWat.dat", "a" ); 
  fprintf( fPtr, "% 3.2f % 3.2f ", siteCoords[ nSites - 1 ].x, siteCoords[ nSites - 1 ].y );
  fprintf( fPtr, "\n" );
  fclose( fPtr );
}

static void PrintSingleLipXY( int lipIndex, FILE *fp )
{
  fprintf( fp, "% 3.2f % 3.2f ", xyLipidComTrue[ lipIndex ].x, xyLipidComTrue[ lipIndex ].y );
  // fflush( fp );
}

void PrintSoluteLatMotion() {
  FILE *fPtr;
  fPtr = fopen( "xySolute.dat", "a" ); 
  fprintf( fPtr, "% 6.4f % 6.4f \n", xySoluteComTrue.x, xySoluteComTrue.y );
  fclose( fPtr ); 
}

void WriteLipLatMotion() 
{
  FILE *fPtr;
  int n, lipCount = 0;
  fPtr = fopen( "xyLipsUpperLayer.dat", "a" ); 
  
  DO_SITE {
    switch ( site[ n ].lipidSpecies ) { 
    case NON_LIPID: break;
    case DOPC:
    case DOPE: // assuming that the DOPC and DOPE models have the same number of CG sites (15)
      if ( siteCoords[ n + 2 ].z > 0 ) // if glycerol site has positive z-coordinate this lipid is in the upper layer
	PrintSingleLipXY( lipCount, fPtr );
      lipCount++;
      n += dopcStruct.nBeads - 1; // set "n" to the index of the next site after this lipid
      break;
    }
  }
  
  if ( lipCount != nLipids ) { printf("*** ERROR in WriteLipLatMotion(): lipCount != nLipids. Exiting...\n"); exit(0); }

  fprintf( fPtr, "\n" );
  fclose( fPtr );
}

void PrintTailOrderParametersOnFile() 
{
  FILE *fPtr;
  if ( ( fPtr = fopen( "tailOrderParameters.dat", "a" ) ) == NULL ) printf( "File tailOrderParameters.dat could not be opened\n" );
  else { 
    fprintf( fPtr, "% .4f % .4f ", ordParTop.sum, ordParMidUnsat.sum ); 
    fprintf( fPtr, "% .4f % .4f ", ordParMidSat.sum, ordParEnd.sum ); 
    fprintf( fPtr, "%14.6f\n",  ordParTail.sum ); // average of the four parameters
    fclose(fPtr);
  }
}

/* FUNCTION PrintSummary produces the output from the run IN PHYSICAL UNITS */
void PrintSummary( FILE *fp )
{
  fprintf( fp, "%8d %9.3f", stepCount, timeNow * TIME_ns );
  if (fp==stdout) {
    fprintf( fp, "% 10.3f", totEnergy.sum );
    fprintf( fp, "% 11.2f", potEnergy.sum ); 
    fprintf( fp, "% 9.1f", pressure.sum * PRESSURE_atm );
    fprintf( fp, "% 8.1f", lipTemperature.sum - 273.15 );
    fprintf( fp, "% 9.1f", watTemperature.sum - 273.15 );
    if ( nSolutes ) fprintf( fp, "%6.1f", solTemperature.sum - 273.15 );
    fprintf( fp, "%12.4f", lipArea.sum * Sqr(LENGTH_Angstrom) ); 
    fprintf( fp, "%13.4f % 11.1f", lipVolume.sum, surfaceTension.sum * TENSION_dyn_cm ); 
    fprintf( fp, "%14.4f % 9.4f", headGroupDipole.sum * DIPOLE_D, headGroupTilt.sum ); 
  } else {
    fprintf( fp, "% 9.3f %5.3f", PropEst( totEnergy ) );
    fprintf( fp, "% 8.2f %4.2f", PropEst( potEnergy ) ); 
    fprintf( fp, "% 7.1f %3.0f ", PropScaleEst( pressure, PRESSURE_atm ) );
    fprintf( fp, "%6.1f %3.1f ",  PropShiftEst( lipTemperature, - 273.15 ) );
    fprintf( fp, "%5.1f %3.1f ",  PropShiftEst( watTemperature, - 273.15 ) );
    if ( nSolutes ) fprintf( fp, "%6.1f %3.1f ", PropShiftEst( solTemperature, - 273.15 ) );
    fprintf( fp, "%8.4f %4.4f", PropScaleEst( lipArea, Sqr(LENGTH_Angstrom) ) ); 
    fprintf( fp, "%8.4f %4.3f %5.1f %2.0f ", PropScaleEst( lipVolume, Cube( LENGTH_nm ) ), PropScaleEst ( surfaceTension, TENSION_dyn_cm ) ); 
    fprintf( fp, "%8.4f %6.4f % 6.4f %6.4f", PropScaleEst( headGroupDipole, DIPOLE_D ), PropEst( headGroupTilt ) ); 
  }
  fprintf( fp, "\n" ); 
  fflush( fp );
}

void PrintSummaryWater( FILE *fp )
{
  fprintf( fp, "%8d %9.3f", stepCount, timeNow *  TIME_ns );
  fprintf( fp, "% 11.4f %7.4f", PropEst( totEnergy ) );
  fprintf( fp, "% 10.2f %7.3f", PropEst( potEnergy ) ); 
  fprintf( fp, "%10.1f %5.0f", PropScaleEst( pressure, PRESSURE_atm ) ); 
  fprintf( fp, "%9.1f %5.1f", PropShiftEst( systemTemperature, - 273.15 ) );
  fprintf( fp, "%9.3f %8.3f", PropScaleEst( dens, BRAHMS_MASS_IN_g / Cube( BRAHMS_LENGTH_IN_cm ) ) ); 
  //  fprintf( fp, "%10.3f", 1000*Sqr( totEnergy.sum2 ) / (kB_IN_kJ_OVER_molK*Sqr( systemTemperature.sum )) ); 
  fprintf( fp, "\n" ); 
  fflush( fp ); 
}

void PrintSummaryOnFile() 
{
  FILE *fPtr;
  if ( ( fPtr = fopen( "summary.dat", "a" ) ) == NULL ) printf( "File summary.dat could not be opened\n" );
  else { PrintSummary( fPtr ); fclose( fPtr ); }
}

void PrintSummaryWaterOnFile() 
{
  FILE *fPtr;
  if ( ( fPtr = fopen( "summary.dat", "a" ) ) == NULL )
    printf( "File summary.dat could not be opened\n" );
  else { 
    PrintSummaryWater( fPtr );
    fclose( fPtr );
  }
}

/************************************************************************************************************
 * UnwrapSitePositionWithRespectToGlycerol -- unwraps coordinates of lipid sites that are wrapped around due 
 * to periodic boundary conditions. This allows each lipid molecule to be visualised "whole" instead of being 
 * potentially split over different periodic images. For visualization purposes only.
 * BUG: if the region is small, this function can occasionally displace a site when there is no need to,
 * causing a visualisation artefact. This bug does not affect the simulation, but it would be good to fix it. 
 ************************************************************************************************************/
static void UnwrapSitePositionWithRespectToGlycerol( VecR *rSiteWrapped, /* out - wrapped position */
						     int n, /* site number */
						     int nSitesToGlycerol ) /* number of sites separating site "n" to the 
									       glycerol site in the same (lipid) molecule */
{
  VecR interSiteDistance;
  VSub( interSiteDistance, siteCoords[ n ], siteCoords[ n-nSitesToGlycerol ] ); 

  VWrapAll( interSiteDistance ); /* wrap distance to nearest image */

  VAdd( *rSiteWrapped, siteCoords[ n-nSitesToGlycerol ], interSiteDistance );
  VScale( *rSiteWrapped, LENGTH_Angstrom );
}

/* PrintPdb -- generates a trajectory file that can be directly read into VMD. */
void PrintPdb( FILE *fPtr )
{
  int n, atomCount = 0;
  VecR tip, tail, rSite, rSiteUnwrapped;
  real dipoleLength = 0.35; 
  /* SSD modelling */
  RMat rMat;
  VecR oh1b, oh2b, /* oxygen-hydrogen vectors in body-fixed frame */
    oob, ohb, t, h1s, os;

  VSet( oh1b, 0.,  0.75,   0.53 ); /* [Angstrom], from Bratko et al, J Chem Phys, 83, 6367, 1985 */
  VSet( oh2b, 0., -0.75,   0.53 );
  VSet( oob, 0.,      0., -0.065 );
  VSet( ohb, 0.,  0.,   0.5 ); /* approximate position of '+' end of water dipole */

  fprintf ( fPtr,"MODEL\n" );
  // Adding unit cell information to display periodic images with VMD - using PDB "CRYST1" record format (to be modified for non-tetragonal cells!):
  fprintf( fPtr, "CRYST1%9.3f%9.3f%9.3f  90.00  90.00  90.00 P  1  1  1    1\n", 
	   region.x*LENGTH_Angstrom, region.y*LENGTH_Angstrom, region.z*LENGTH_Angstrom );
  DO_SITE {
    rMat = siteRotMat[ n ];      
    VSCopy( rSite, LENGTH_Angstrom, siteCoords[n] );
    switch ( site[ n ].lipidUnit ) {
    case WATER: /* SSD sites represented as 2-site molecules, to highlight the dipole: "H" for the "+" end and "O" for the "-" */
      MVMulT( t, rMat.u, oob );
      VAdd( os, rSite, t );
      fprintf( fPtr, "ATOM %6d %4s WAT %5d %11.3f%8.3f%8.3f\n", ++atomCount, "OW", n+1, os.x, os.y, os.z );
      MVMulT( t, rMat.u, ohb );
      VAdd( h1s, rSite, t );
      fprintf( fPtr, "ATOM %6d %4s WAT %5d %11.3f%8.3f%8.3f\n", ++atomCount, "HW", n+1, h1s.x, h1s.y, h1s.z );
      break;
    case CHOLINE_OR_AMINE:
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, -2 );
      if ( site[n].type == CHOLINE_TYPE )
	fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "NC3", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      else if ( site[n].type == AMINE_TYPE )
	fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "NH3", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      break;
    case PHOSPHATE:
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, -1 );
      fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "PO4", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      break;
    case GLYCEROL:
      /*       fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "C", n+1,  rSite.x,  rSite.y,  rSite.z ); */
      VSAdd( tip,  rSite,  dipoleLength, siteOrientVec[ n ] );
      fprintf ( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "C", n+1, tip.x, tip.y, tip.z );
      VSAdd( tail, rSite, -dipoleLength, siteOrientVec[ n ] );
      fprintf ( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "O", n+1, tail.x, tail.y, tail.z );
      break;
    case ESTER_A:  
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, 1 );
      /* fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "O", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z ); */
      VSAdd( tip, rSiteUnwrapped,  dipoleLength, siteOrientVec[ n ] );
      fprintf ( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "O", n+1, tip.x, tip.y, tip.z );
      VSAdd( tail,rSiteUnwrapped, -dipoleLength, siteOrientVec[ n ] );
      fprintf ( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "C", n+1, tail.x, tail.y, tail.z );
      break;
    case ESTER_B: /* ester sites represented as 2-site molecules, to highlight the dipole: "H" for the "+" end and "O" for the "-" */
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, 7 );
      /* fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "O", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z ); */
      VSAdd( tip, rSiteUnwrapped,  dipoleLength, siteOrientVec[ n ] );
      fprintf ( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "O", n+1, tip.x, tip.y, tip.z );
      VSAdd( tail,rSiteUnwrapped, -dipoleLength, siteOrientVec[ n ] );
      fprintf ( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "C", n+1, tail.x, tail.y, tail.z );
      break;
    case TAIL_A1:  
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, 2 );
      fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "CH2", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      break;
    case TAIL_A2:  
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, 3 );
      fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "CH2", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      break;
    case TAIL_A3:
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, 4 );
      fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "CH2", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      break;
    case TAIL_A4: 
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, 5 );
      fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "CH2", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      break;
    case TAIL_A5:  
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, 6 );
      fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "CH3", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      break;
    case TAIL_B1:  
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, 8 );
      fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "CH2", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      break;
    case TAIL_B2:  
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, 9 );
      fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "CH2", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      break;
    case TAIL_B3:  
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, 10 );
      fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "CH2", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      break;
    case TAIL_B4:  
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, 11 );
      fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "CH2", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      break;
    case TAIL_B5:  
      UnwrapSitePositionWithRespectToGlycerol( &rSiteUnwrapped, n, 12 );
      fprintf( fPtr, "ATOM %6d %4s UNK %5d %11.3f%8.3f%8.3f\n", ++atomCount, "CH3", n+1, rSiteUnwrapped.x, rSiteUnwrapped.y, rSiteUnwrapped.z );
      break;
    default: printf("ERROR in PrintPdb switching!!!\n"); exit(0);
    } /* end switch */
  } /* end DO_SITE */
  fprintf ( fPtr,"ENDMDL\n" );
}

void PutInitialConfiguration() 
{
  FILE *fPtr;

  fPtr = fopen( "initialConfiguration.pdb", "w" ); 
  PrintPdb( fPtr ); 
  fclose( fPtr );

  fPtr = fopen( "trajectory.pdb", "a" ); 
  PrintPdb( fPtr ); 
  fclose( fPtr );
}

void WriteAreaVolume() 
{
  fprintf( fPtrAV, "%.3f %.2f\n", xyArea.val, boxVolume.val ); /* [nm^2], [nm^3] */
}
